import pickle
import numpy as np
import scipy.io.wavfile as wf
import python_speech_features as sf
from flask import Flask, request, jsonify
from flask_cors import CORS
import os
from pydub import AudioSegment
import logging

# 设置日志记录
logging.basicConfig(level=logging.INFO)

class AudioRecognizer:
    def __init__(self, model_path):
        self.names = ['abnormal_treble', 'normal_bass', 'normal_bass_inside_the_carriage', 'normal_treble', 'normal_treble_inside_the_carriage', 'resonant_bass_inside_the_carriage', 'resonant_treble_inside_the_carriage']
        with open(model_path, 'rb') as file:
            self.model = pickle.load(file)

    def mfcc(self, file):
        try:
            sample_rate, signs = wf.read(file)
            mfc = sf.mfcc(signs, sample_rate)
            sample = np.mean(mfc, axis=0)
            return sample.reshape(1, -1)
        except ValueError as e:
            logging.error(f"Error reading the WAV file: {e}")
            return None

    def predict(self, file_path):
        features = self.mfcc(file_path)
        if features is not None:
            prediction = self.model.predict(features)
            return self.names[prediction[0]]
        else:
            return "Error processing audio file"

def convert_to_wav(file_path):
    # 检测文件格式并转换为WAV
    if not file_path.lower().endswith('.wav'):
        audio = AudioSegment.from_file(file_path)
        new_file_path = file_path.rsplit('.', 1)[0] + '.wav'
        audio.export(new_file_path, format='wav')
        os.remove(file_path)  # 删除原始文件
        return new_file_path
    return file_path

app = Flask(__name__)
CORS(app)

recognizer = AudioRecognizer(model_path="model.pickle")
ALLOWED_EXTENSIONS = set(['aac', 'wav', 'mp3'])
UPLOAD_FOLDER = 'uploads'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

os.makedirs(UPLOAD_FOLDER, exist_ok=True)

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/upload', methods=['POST'])
def upload_audio():
    def upload_audio():
        if 'file' not in request.files:
            return jsonify(status='error', message='No file part'), 400
        file = request.files['file']

        if file.filename == '':
            return jsonify(status='error', message='No selected file'), 400
        if file and allowed_file(file.filename):
            filename = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
            file.save(filename)

            predicted_label = recognizer.predict(filename)
            return jsonify(status='success', result=predicted_label), 200
        else:
            return jsonify(status='error', message='File type not allowed'), 400

@app.route('/predict', methods=['POST'])
def predict_audio():
    data = request.json
    file_path = data.get('file_path')
    if not file_path or not os.path.isfile(file_path):
        return jsonify(status='error', message='File not found'), 400
    if not allowed_file(file_path):
        return jsonify(status='error', message='File type not allowed'), 400

    # 转换为WAV格式
    wav_filename = convert_to_wav(file_path)

    predicted_label = recognizer.predict(wav_filename)
    if predicted_label == "Error processing audio file":
        return jsonify(status='error', message='Error processing audio file'), 500
    return jsonify(status='success', result=predicted_label), 200

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False)